import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

from verify_code.ctc_model import DigitsModel, myctc_loss
from verify_code.main import CodeDataset
import numpy as np

if __name__ == '__main__':
    epochs = 2
    batch_size = 4

    my_transforms = transforms.Compose([
        transforms.ToPILImage(),  # 不转换为PIL会报错
        transforms.Resize([50, 50]),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5,), std=(0.5, 0.5, 0.5,))
    ])
    dataset = CodeDataset(path=r'C:\Users\Administrator\Desktop\WangTing\five_numbers_verification_code',
                          transform=my_transforms)
    t_dataset = CodeDataset(path=r'C:\Users\Administrator\Desktop\WangTing\already_train',
                          transform=my_transforms)  # , transform=transforms.ToTensor()

    # t_dataset = CodeDataset(path=r'C:\Users\Administrator\Downloads\qingdao',
    #                       transform=my_transforms)  # 不能直接应用于新的验证码
    test_dataloader = torch.utils.data.DataLoader(t_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

    label_map = dataset.characters
    # DIGITS_MAP = dataset.map

    char_kinds = dataset.myclass_len

    mymodel = DigitsModel(char_kinds)
    mymodel.load_state_dict(torch.load(r"verify_code/models/checkpoints.pt"))
    mymodel.eval()
    # print(mymodel)

    ctc = myctc_loss()
    for data_batch, labels_batch in test_dataloader:
        logits = mymodel(data_batch)
        x = torch.transpose(logits, 0, 1)
        # # x = torch.transpose(x, 2, 1)
        probs = F.softmax(x, dim=2).data.cpu()

        from fast_ctc_decode import beam_search, viterbi_search
        for d,l in zip(probs,labels_batch):
            seq, path = beam_search(d.detach().numpy(), label_map, beam_size=3,beam_cut_threshold=0.0001)
            print([dataset.num_to_char[i] for i in np.array(l).tolist()])
            print(seq)
            pass
